import os
import platform
import sys

import torch

# 添加当前目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
print("sys.path:", sys.path)
print(sys.version)
print(platform.architecture)
print(torch.__version__, sys.version)


# 加载自定义算子库
lib_path = os.path.join(os.path.dirname(__file__), "dose_torch.pyd")
if os.path.exists(lib_path):
    torch.ops.load_library(lib_path)
else:
    # 如果库不存在, 尝试从当前目录加载
    lib_path = os.path.join(os.path.dirname(__file__), "libdose_stat_torch.so")
    if os.path.exists(lib_path):
        torch.ops.load_library(lib_path)

print("from dose_torch import *")
